import csv
import os
import torch
import numpy as np
from skimage import io, transform
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms, utils

class TraficSignDataSet(Dataset):

    def __init__(self, need_train, transform=None):
        if need_train is True:
            self.img_dir = './Final_Training_ROI/'
            self.label_dir = './Final_Training_Label/'
        else:
            self.img_dir = './Final_Test_ROI/'
            self.label_dir = './Final_Test_Label/'
        self.csv_filename = self.label_dir + 'label.csv'
        self.img_filename_list = []
        self.img_label_list = []
        self.transform = transform

        # Extract image filenames and labels from the extracted csv file
        with open(self.csv_filename) as f:
            reader = csv.reader(f)
            for row in reader:  
                img_filename, label = row[0], row[1]
                self.img_filename_list.append(img_filename)
                self.img_label_list.append(int(label))
       

    def __len__(self):
        return len(self.img_filename_list)

    def __getitem__(self, idx):
        img = io.imread(self.img_dir + self.img_filename_list[idx])
        label = self.img_label_list[idx]
        
        if self.transform:
            img = self.transform(img)

        sample = {'image': img, 'label': label}
        return sample

